import torch
import torch.nn as nn
from .yolov1_basic import Conv


class SPPF(nn.Module):
    def __init__(
        self,
        in_dim,
        out_dim,
        expand_ratio=0.5,
        pooling_size=5,
        act_type="lrelu",
        norm_type="bn",
    ):
        super(SPPF, self).__init__()
        intermediate_dim = int(in_dim * expand_ratio)
        self.cv1 = Conv(in_dim, intermediate_dim, 1, 1, 0, 1, act_type, norm_type)
        self.cv2 = Conv(intermediate_dim * 4, out_dim, 1, 1, 0, 1, act_type, norm_type)
        self.m = nn.MaxPool2d(
            kernel_size=pooling_size, stride=1, padding=pooling_size // 2
        )
        self.out_dim = out_dim

    def forward(self, x):
        x = self.cv1(x)
        y1 = self.m(x)
        y2 = self.m(y1)
        y3 = self.m(y2)
        y = self.cv2(torch.concat([x, y1, y2, y3], dim=1))
        return y


def build_neck(cfg, in_dim, out_dim):
    model = cfg["neck"]
    print("====================")
    print("model neck:{}".format(model))
    if model == "sppf":
        return SPPF(
            in_dim,
            out_dim,
            cfg["expand_ratio"],
            cfg["pooling_size"],
            cfg["neck_act"],
            cfg["neck_norm"],
        )
    else:
        raise NotImplementedError("neck {} not implemented".format(model))
